import os
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
import json
import random
import time
from typing import Union
from fastapi import FastAPI
import paho.mqtt.client as mqtt
from paho.mqtt import publish


from modelscope import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

model_dir = snapshot_download('qwen/Qwen-7B-Chat',cache_dir='./check')


tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",
    trust_remote_code=True,

).eval()



# The callback for when the client receives a CONNACK response from the server.
def on_connect(client, userdata, flags, reason_code, properties):
    print(f"Connected with result code {reason_code}")
    # Subscribing in on_connect() means that if we lose the connection and
    # reconnect then subscriptions will be renewed.
    client.subscribe("llm/ask")

# The callback for when a PUBLISH message is received from the server.
def on_message(client, userdata, msg):
    msg = json.loads(msg.payload)
    print(msg)
    if len(msg['history']) ==0:
        msg, history = model.chat(tokenizer, msg['message'], history=None)
    else:
        msg, history = model.chat(tokenizer, msg['message'], history=msg['history'])

    publish.single("llm/answer", json.dumps({'message':msg, 'history':history}), hostname="120.77.8.146")

mqttc = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
mqttc.on_connect = on_connect
mqttc.on_message = on_message

mqttc.connect("120.77.8.146", 1883, 60)
mqttc.loop_forever()

